{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "import os\n",
    "from shared.utils import load_images, plot_images_with_max_per_row, plot_attn_maps, filter_layers\n",
    "from shared.algorithms import find_register_neurons\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\n",
    "import yaml\n",
    "from tqdm import tqdm\n",
    "\n",
    "sys.path.append(\"clip/\")\n",
    "sys.path.append(\"dinov2/\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set cuda visible devices\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\" # Change this to the GPU you want to use"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL = \"dinov2\" # can be \"clip\" or \"dinov2\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if MODEL == \"dinov2\":\n",
    "  from dinov2_state import load_dinov2_state\n",
    "\n",
    "  with open(\"configs/dinov2_large.yaml\", \"r\") as f:\n",
    "    config = yaml.safe_load(f)\n",
    "\n",
    "  state = load_dinov2_state(config)\n",
    "elif MODEL == \"clip\":\n",
    "  from clip_state import load_clip_state\n",
    "\n",
    "  with open(\"configs/openclip_base.yaml\", \"r\") as f:\n",
    "    config = yaml.safe_load(f)\n",
    "\n",
    "  state = load_clip_state(config)\n",
    "else:\n",
    "  raise ValueError(f\"Model {MODEL} not supported\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGENET_PATH = \"/datasets/ilsvrc/current/val\" # Pass in path to ImageNet\n",
    "IMAGE_SIZE = 224 # Preprocessed image size\n",
    "\n",
    "run_model = state[\"run_model\"]\n",
    "model = state[\"model\"]\n",
    "preprocess = state[\"preprocess\"] # Preprocess function for input images\n",
    "hook_manager = state[\"hook_manager\"]\n",
    "num_layers = state[\"num_layers\"]\n",
    "num_heads = state[\"num_heads\"]\n",
    "patch_size = state[\"patch_size\"]\n",
    "config = state[\"config\"]\n",
    "patch_height = IMAGE_SIZE // patch_size\n",
    "patch_width = IMAGE_SIZE // patch_size\n",
    "device = \"cuda:0\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Run on one image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Load an image for analysis on the original model (no test-time register). There should be artifacts appearing in uniform regions (ex. background) in the attention map. There should also be high-norm outliers in the patch norms in the latter layers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = load_images(IMAGENET_PATH, count = 1)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "processed_image = preprocess(image).unsqueeze(0).to(device)\n",
    "hook_manager.reinit()\n",
    "hook_manager.finalize()\n",
    "representation = run_model(model, processed_image)\n",
    "\n",
    "attention_maps = hook_manager.get_attention_maps() # shape (L, H, N, N)\n",
    "layer_outputs = hook_manager.get_layer_outputs() # shape (L, N, D)\n",
    "\n",
    "patch_norms = np.linalg.norm(layer_outputs[:, 1:], axis=2).reshape(num_layers, patch_height, patch_width)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(image)\n",
    "plt.axis('off')\n",
    "plt.title(\"Input\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Plot patch norms across all layers. Notice that outliers appear in the later layers.\n",
    "\n",
    "plt = plot_images_with_max_per_row(patch_norms, max_per_row = 4, image_title = \"Layer\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot attention maps\n",
    "\n",
    "cls_attn_maps = attention_maps[:, :, 0, 1:].reshape(num_layers, num_heads, patch_height, patch_width)\n",
    "plt = plot_attn_maps(cls_attn_maps)\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Identify where outliers & attention sinks appear"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Look for the layer(s) where the outliers begin appearing. The attention sinks usually show up at the layer after."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rand_images = load_images(IMAGENET_PATH, count = 10)\n",
    "max_patch_norms = [0] * num_layers\n",
    "max_attn_norms = [0] * num_layers\n",
    "\n",
    "for image in tqdm(rand_images, desc=\"Processing images\"):\n",
    "  processed_image = preprocess(image).unsqueeze(0).to(device)\n",
    "  hook_manager.reinit()\n",
    "  hook_manager.finalize()\n",
    "  representation = run_model(model, processed_image)\n",
    "  attention_maps = hook_manager.get_attention_maps() # shape (L, H, N, N)\n",
    "  layer_outputs = hook_manager.get_layer_outputs() # shape (L, N, D)\n",
    "\n",
    "  patch_norms = np.max(np.linalg.norm(layer_outputs[:, 1:patch_height * patch_width + 1], axis=2), axis=1)\n",
    "  attn_norms = np.max(np.mean(attention_maps[:, :, 0, 1:patch_height * patch_width + 1], axis = 1), axis = 1)\n",
    "\n",
    "  for j in range(num_layers):\n",
    "    max_attn_norms[j] += attn_norms[j]\n",
    "    max_patch_norms[j] += patch_norms[j]\n",
    "\n",
    "max_attn_norms = [x / len(rand_images) for x in max_attn_norms]\n",
    "max_patch_norms = [x / len(rand_images) for x in max_patch_norms]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Create a figure with two subplots side by side\n",
    "fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(18, 5), dpi=100)\n",
    "\n",
    "# === Subplot 1: Patch Norms ===\n",
    "ax1.plot(range(num_layers), max_patch_norms, marker='o', markersize=8, color='steelblue', linestyle='-', linewidth=2)\n",
    "ax1.set_title('Average Max Patch Norms (Layer output)')\n",
    "ax1.set_xlabel('Layer')\n",
    "ax1.set_ylabel('Norm')\n",
    "ax1.set_xticks(range(0, num_layers, 2))  # Show every second tick\n",
    "ax1.tick_params(axis='both')\n",
    "ax1.tick_params(axis='x', which='major', pad=15)\n",
    "ax1.margins(x=0.1)\n",
    "ax1.grid(True, linestyle='-')\n",
    "\n",
    "# === Subplot 2: Attention Norms ===\n",
    "ax2.plot(range(num_layers), max_attn_norms, marker='^', markersize=8, color='orange', linestyle='-', linewidth=2)\n",
    "ax2.set_title('Average Max Attention (CLS)')\n",
    "ax2.set_xlabel('Layer')\n",
    "ax2.set_ylabel('Attention Value')\n",
    "ax2.set_xticks(range(0, num_layers, 2))  # Show every second tick\n",
    "ax2.tick_params(axis='both')\n",
    "ax2.tick_params(axis='x', which='major', pad=15)\n",
    "ax2.margins(x=0.1)\n",
    "ax2.grid(True, linestyle='-')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Identify register neurons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "#########################################\n",
    "#              PARAMETERS               #\n",
    "#########################################\n",
    "\n",
    "# layer used to detect outliers based on the patch norms. Set to the last layer (-1) for CLIP and second-to-last layer (-2) for DINOv2 large\n",
    "detect_outliers_layer = config[\"detect_outliers_layer\"]\n",
    "register_norm_threshold = config[\"register_norm_threshold\"] # threshold for detecting register neurons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "register_neurons = find_register_neurons(\n",
    "    model_state=state,\n",
    "    image_path=IMAGENET_PATH,\n",
    "    detect_outliers_layer=detect_outliers_layer,\n",
    "    processed_image_cnt=100,\n",
    "    register_norm_threshold=register_norm_threshold,\n",
    "    apply_sparsity_filter=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Optionally, you can save the register neurons"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(register_neurons, \"register_neurons.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "register_neurons = torch.load(\"register_neurons.pt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Evaluate register neurons w/ test-time registers"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We want to find the right parameters for the register neuron intervention such that the test-time register can absorb the outliers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "#########################################\n",
    "#              PARAMETERS               #\n",
    "#########################################\n",
    "top_k = config[\"top_k\"]\n",
    "highest_layer = config[\"highest_layer\"]\n",
    "num_registers = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "filtered_register_neurons = filter_layers(register_neurons, highest_layer = highest_layer)\n",
    "\n",
    "neurons_to_ablate = dict()\n",
    "for (layer, neuron, score) in filtered_register_neurons[:top_k]:\n",
    "  if layer not in neurons_to_ablate:\n",
    "    neurons_to_ablate[layer] = []\n",
    "  neurons_to_ablate[layer].append(neuron)\n",
    "print(neurons_to_ablate)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "random_images = load_images(IMAGENET_PATH, count = 50)\n",
    "image_norms = []\n",
    "register_norms = []\n",
    "image_attentions = []\n",
    "register_attentions = []\n",
    "for i in tqdm(range(len(random_images)), desc = \"Processing random images\"):\n",
    "  image = preprocess(random_images[i]).unsqueeze(0).to(device)\n",
    "\n",
    "  hook_manager.reinit()\n",
    "  hook_manager.intervene_register_neurons(neurons_to_ablate=neurons_to_ablate, num_registers = num_registers, normal_values=\"zero\", scale = 1)\n",
    "  hook_manager.finalize()\n",
    "  representation = run_model(model, image, num_registers = num_registers)\n",
    "  attention_maps = hook_manager.get_attention_maps()\n",
    "  layer_outputs = hook_manager.get_layer_outputs()\n",
    "\n",
    "  layer_norms = np.linalg.norm(layer_outputs[detect_outliers_layer], axis=1)\n",
    "\n",
    "  image_patch_norms = layer_norms[1:patch_height * patch_width + 1]\n",
    "  register_patch_norms = layer_norms[patch_height * patch_width + 1:]\n",
    "\n",
    "  image_norms.extend(image_patch_norms.tolist())\n",
    "  register_norms.extend(register_patch_norms.tolist())\n",
    "\n",
    "  image_attentions.append(np.max(attention_maps[-1, :, 0, 1:patch_height * patch_width + 1]))\n",
    "  register_attentions.append(np.max(attention_maps[-1, :, 0, patch_height * patch_width + 1:]))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Plot the distribution of norms for the image patches and extra register token. We also track the max attention among the last layer's heads that the CLS token attends to the image patches and register token. If the parameters are well-set, we should see that the distributions of norms and attention for the register token should be much higher than the image patches. Furthermore, the norms of the image patches should be roughly normal and not show a skew (which means there are still outliers).\n",
    "\n",
    "Tips if you don't observe this:\n",
    "1. Increase top_k. However, the fewer neurons we ablate, the better.\n",
    "2. Increase the highest_layer we're filtering the register neurons below. Start by setting it to the layer at which outliers appear.\n",
    "3. Increase the lowest layer we filter above (pass `lowest_layer` into `filter_layers`). Generally speaking, we don't want to ablate neurons too early."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the image norms and register norms as histograms side by side\n",
    "plt.figure(figsize=(10, 12))\n",
    "\n",
    "# Plot image norms\n",
    "plt.subplot(2, 2, 1)\n",
    "plt.hist(image_norms, bins=50, alpha=0.7, color='blue')\n",
    "plt.xlabel('Norm Value')\n",
    "plt.ylabel('Frequency')\n",
    "plt.title('Image Patch Norms')\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.axvline(x=np.median(image_norms), color='r', linestyle='--',\n",
    "            label=f'Median: {np.median(image_norms):.2f}')\n",
    "plt.legend()\n",
    "\n",
    "# Plot register norms\n",
    "plt.subplot(2, 2, 2)\n",
    "plt.hist(register_norms, bins=50, alpha=0.7, color='green')\n",
    "plt.xlabel('Norm Value')\n",
    "plt.ylabel('Frequency')\n",
    "plt.title('Register Patch Norms')\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.axvline(x=np.median(register_norms), color='r', linestyle='--',\n",
    "            label=f'Median: {np.median(register_norms):.2f}')\n",
    "plt.legend()\n",
    "\n",
    "# Plot image attentions\n",
    "plt.subplot(2, 2, 3)\n",
    "plt.hist(image_attentions, bins=50, alpha=0.7, color='blue')\n",
    "plt.xlabel('Attention Value')\n",
    "plt.ylabel('Frequency')\n",
    "plt.title('Image Patch Attentions')\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.axvline(x=np.median(image_attentions), color='r', linestyle='--',\n",
    "            label=f'Median: {np.median(image_attentions):.2f}')\n",
    "plt.legend()\n",
    "\n",
    "# Plot register attentions\n",
    "plt.subplot(2, 2, 4)\n",
    "plt.hist(register_attentions, bins=50, alpha=0.7, color='green')\n",
    "plt.xlabel('Attention Value')\n",
    "plt.ylabel('Frequency')\n",
    "plt.title('Register Patch Attentions')\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.axvline(x=np.median(register_attentions), color='r', linestyle='--',\n",
    "            label=f'Median: {np.median(register_attentions):.2f}')\n",
    "plt.legend()\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()\n",
    "\n",
    "# Print some statistics for comparison\n",
    "print(f\"Image norms - Min: {min(image_norms):.2f}, Max: {max(image_norms):.2f}, Mean: {np.mean(image_norms):.2f}\")\n",
    "print(f\"Register norms - Min: {min(register_norms):.2f}, Max: {max(register_norms):.2f}, Mean: {np.mean(register_norms):.2f}\")\n",
    "print(f\"Image attentions - Min: {min(image_attentions):.2f}, Max: {max(image_attentions):.2f}, Mean: {np.mean(image_attentions):.2f}\")\n",
    "print(f\"Register attentions - Min: {min(register_attentions):.2f}, Max: {max(register_attentions):.2f}, Mean: {np.mean(register_attentions):.2f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load an image to check for no outliers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = load_images(IMAGENET_PATH, count = 1)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "plt.imshow(image)\n",
    "plt.axis('off')\n",
    "plt.title(\"Input\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "processed_image = preprocess(image).unsqueeze(0).to(device)\n",
    "\n",
    "hook_manager.reinit()\n",
    "hook_manager.finalize()\n",
    "representation = run_model(model, processed_image)\n",
    "original_attention_maps = hook_manager.get_attention_maps()\n",
    "original_layer_outputs = hook_manager.get_layer_outputs()\n",
    "\n",
    "hook_manager.reinit()\n",
    "hook_manager.intervene_register_neurons(neurons_to_ablate=neurons_to_ablate, num_registers = num_registers)\n",
    "hook_manager.finalize()\n",
    "representation = run_model(model, processed_image, num_registers = num_registers)\n",
    "ablated_attention_maps = hook_manager.get_attention_maps()\n",
    "ablated_layer_outputs = hook_manager.get_layer_outputs()\n",
    "ablated_neuron_activations = hook_manager.get_neuron_activations()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Norm map of output patch embeddings - baseline and ablated comparison for all layers\n",
    "\n",
    "# Import the necessary module for make_axes_locatable\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "\n",
    "# Create a figure with subplots for each layer (1 row per layer)\n",
    "fig, axs = plt.subplots(num_layers, 2, figsize=(16, 4 * num_layers))\n",
    "fig.subplots_adjust(hspace=0.5, wspace=0.3)\n",
    "\n",
    "# Plot norm maps for each layer\n",
    "for layer in range(num_layers):\n",
    "    # Calculate norms for baseline and ablated outputs\n",
    "\n",
    "    # Calculate norms for outputs\n",
    "    baseline_output_norms_flat = np.linalg.norm(original_layer_outputs[layer, 1:], axis=1)\n",
    "    ablated_output_norms_flat = np.linalg.norm(ablated_layer_outputs[layer, 1:], axis=1)\n",
    "\n",
    "    # Handle non-square reshaping\n",
    "    def reshape_with_extras(flat_array, patch_height, patch_width):\n",
    "        total_patches = len(flat_array)\n",
    "        if total_patches == patch_height * patch_width:\n",
    "            # Perfect square case\n",
    "            return flat_array.reshape((patch_height, patch_width)), None\n",
    "        else:\n",
    "            # Non-square case\n",
    "            square_part = flat_array[:patch_height * patch_width].reshape((patch_height, patch_width))\n",
    "            extra_part = flat_array[patch_height * patch_width:]\n",
    "            return square_part, extra_part\n",
    "\n",
    "    # Reshape with handling for extra values\n",
    "    baseline_output_norms, baseline_output_extras = reshape_with_extras(baseline_output_norms_flat, patch_height, patch_width)\n",
    "    ablated_output_norms, ablated_output_extras = reshape_with_extras(ablated_output_norms_flat, patch_height, patch_width)\n",
    "\n",
    "    # Plot baseline output\n",
    "    im3 = axs[layer, 0].imshow(baseline_output_norms, cmap='viridis')\n",
    "    extra_info = \"\"\n",
    "    axs[layer, 0].set_title(f'Layer {layer} - Original (output){extra_info}', fontsize=14)\n",
    "    axs[layer, 0].set_xlabel('Patch X', fontsize=12)\n",
    "    axs[layer, 0].set_ylabel('Patch Y', fontsize=12)\n",
    "\n",
    "    # Add colorbar for output original\n",
    "    divider = make_axes_locatable(axs[layer, 0])\n",
    "    cax = divider.append_axes(\"right\", size=\"5%\", pad=0.1)\n",
    "    cbar = fig.colorbar(im3, cax=cax)\n",
    "    cbar.set_label(f'Layer {layer} Output Norm', fontsize=12)\n",
    "\n",
    "    # Plot ablated output\n",
    "    im4 = axs[layer, 1].imshow(ablated_output_norms, cmap='viridis')\n",
    "    extra_info = \"\"\n",
    "    axs[layer, 1].set_title(f'Layer {layer} - Ablated (output){extra_info}', fontsize=14)\n",
    "    axs[layer, 1].set_xlabel('Patch X', fontsize=12)\n",
    "    axs[layer, 1].set_ylabel('Patch Y', fontsize=12)\n",
    "\n",
    "    # Add colorbar for output ablated\n",
    "    divider = make_axes_locatable(axs[layer, 1])\n",
    "    cax = divider.append_axes(\"right\", size=\"5%\", pad=0.1)\n",
    "    cbar = fig.colorbar(im4, cax=cax)\n",
    "    cbar.set_label(f'Layer {layer} Output Norm', fontsize=12)\n",
    "\n",
    "fig.suptitle('Norm Maps of Image Patches Across All Layers', fontsize=20, y=1.00)\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Original attention maps\n",
    "plt = plot_attn_maps(original_attention_maps[:, :, 0, 1:patch_height * patch_width + 1].reshape((num_layers, num_heads, patch_height, patch_width)))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Ablated attention maps\n",
    "plt = plot_attn_maps(ablated_attention_maps[:, :, 0, 1:patch_height * patch_width + 1].reshape((num_layers, num_heads, patch_height, patch_width)))\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vit_register_neurons3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
